import torch.nn as nn
from torch._export import capture_pre_autograd_graph
from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_qat_pt2e

from coremltools.optimize.torch.quantization._coreml_quantizer import CoreMLQuantizer

model = nn.Sequential(
    OrderedDict(
        {
            "conv": nn.Conv2d(1, 20, (3, 3)),
            "relu1": nn.ReLU(),
            "conv2": nn.Conv2d(20, 20, (3, 3)),
            "relu2": nn.ReLU(),
        }
    )
)

loss_fn = define_loss()

# initialize the annotator with quantization config
config = LinearQuantizerConfig.from_dict(
    {
        "global_config": {
            "quantization_scheme": "symmetric",
        }
    }
)
quantizer = CoreMLQuantizer(config)

example_inputs = torch.randn(1, 1, 28, 28)

# create export graph
exported_model = capture_pre_autograd_graph(model, (example_inputs,))

# prepare the model to insert FakeQuantize layers for QAT
prepared_model = prepare_qat_pt2e(exported_model, quantizer)

# use prepared model in your PyTorch training loop
for inputs, labels in data:
    output = prepared_model(inputs)
    loss = loss_fn(output, labels)
    loss.backward()
    optimizer.step()
    # turn observers/quantizers on/off depending on iteration number

# convert operations to their quanitzed counterparts using parameters learnt via QAT
model = convert_pt2e(prepared_model)